
import numpy as np


def viterbi_algo(text, transition, emission, ind2tag, word2ind):
    """
    pos tagger的viterbi算法

    Args:
        text(iterables)      :- 要估计的句子的序列,长度为l
        transition(np.array) :- 词性转换矩阵，
                                [i,j]元素表示从i词性到j词性的概率，
                                维度为[K ,K], K为词性的种类数
        emission(np.array)   :- 产生词的概率矩阵,
                                [i,k]元素表示i词性生成k词的概率,
                                维度为[K, v], v为字典的大小
        ind2tag(iterables)   :- 第i个元素为i词性
        word2ind(dict)       :- k词为key，序号为value
    Returns:
        paths(np.array)      :- [i,w]元素表示第w词为i词性时上一个词的词性，
                                维度为[n, l]
        viterbi(np.array)    :- [i,w]元素表示第w词为i词性的概率,
                                维度为[K,l]

    """
    try:
        assert transition.shape[0] == transition.shape[1]
    except AssertionError:
        print("转移矩阵不是方阵")

    try:
        assert transition.shape[0] == emission.shape[0]
    except AssertionError:
        print("emission矩阵的行数和词性数不一致")

    try:
        assert len(ind2tag) == transition.shape[0]
    except AssertionError:
        print("ind2tag长度和词性数不一致")

    cols = len(text)
    rows = transition.shape[0]

    paths = np.zeros((rows, cols))
    viterbi = np.zeros((rows, cols))
    viterbi[0, 0] = 1
    for j in range(1, cols):
        for i in range(rows):
            prob = viterbi[:, j - 1] * transition[:, i] * emission[i, word2ind[text[j]]]  # [cols, 1]
            sort = np.argsort(prob)
            paths[i, j] = sort[-1]
            viterbi[i, j] = max(prob)

    last = int(np.argsort(viterbi[:, j])[-1])  # 最后一个词的概率最大的行序
    print(f"词为{text[-1]}, 词性为{ind2tag[last]}, 概率为{viterbi[last, j] * 100}%")

    for j in range(cols - 1):
        last = int(paths[last, cols - 1 - j])  # 上一个词的行序
        print(f"词为{text[-2 - j]},词性为{ind2tag[last]},概率为{viterbi[last, cols - 2 - j] * 100}%")

    return paths, viterbi


def demo():
    text = ['b', 'I', 'love', 'dogs', 'e']

    ind2tag = ['b', 'rr', 'v', 'n', 'e']  # 开头，人称代词，动词，名词

    transition = np.array(
        [
            [0, 0.4, 0.2, 0.4, 0],  # 从b到 rr, verb, noun, e
            [0, 0, 0.85, 0.1, 0.05],
            [0, 0.3, 0, 0.6, 0.1],
            [0, 0, 0.6, 0, 0.4],
            [0, 0, 0, 0, 0]
        ]
    )

    emission = np.array(
        [
            [1, 0, 0, 0, 0],
            [0, 0.4, 0, 0.6, 0],
            [0, 0.45, 0.55, 0, 0],
            [0, 0, 0, 1, 0],
            [0, 0, 0, 0, 1]
        ]

    )

    word2ind = {
        'b': 0,
        'I': 1,
        'love': 2,
        'dogs': 3,
        'e': 4
    }
    paths, viterbi = viterbi_algo(text, transition, emission, ind2tag, word2ind)
    print(viterbi)
    print(paths)


if __name__ == "__main__":
    demo()